import sys
sys.path.append('..')
import os
import glob
import pandas as pd
import numpy as np
import json
import argparse
from tqdm import tqdm
from utils.helper import get_subdirs
from pathlib import Path
from PIL import Image
from sklearn.model_selection import train_test_split
import random
import cv2
from typing import Dict, List
from utils.json import json_to_image

def parse_args():
    """
    Set args parameters
    """
    parser = argparse.ArgumentParser(description='Split data for training Seg model.')
    parser.add_argument("--json_path", type=str, default='./path_small.json', help="Define the ratio of validation set.")
    parser.add_argument("--save_dir", type=str, default='/root/autodl-tmp/adc_T10/270H1_T270/train_data/SMALL_IMGS', help="Define the dir to store csv files.")
    parser.add_argument("--no_defect_code", nargs='+', default=['TSFAS', 'TSDFS', 'NOCOD', 'NOCOD2', 'TFOL0', 'TSFIX'], help="Define the dir to store csv files.")
    parser.add_argument("--img_ext", nargs='+', default=['JPG', 'jpg'], help="img format for training samples.")
    parser.add_argument("--small_code", nargs='+', default=['TS15', 'TEDP'], help="Define the dir to store csv files.")
    parser.add_argument("--grid_num", type=int, default=2, help="Define grid_num*grid_num patches .")
    args = parser.parse_args()

    return args

def get_extensions(path):
    ListFiles = os.walk(path)
    SplitTypes = []
    for walk_output in ListFiles:
        for file_name in walk_output[-1]:
            SplitTypes.append(file_name.split(".")[-1])
    extensions, counts = np.unique(SplitTypes, return_counts=True)
    return extensions[np.argmax(counts)]

def check_json(json_path):
    with open(json_path, 'r+', encoding="utf-8") as f:
        data = json.load(f)
    label_check = True
    for i in range(len(data['shapes'])):
        shape_slice = data['shapes'][i]
        points_slice = shape_slice['points']
        if len(points_slice) < 3:
            label_check = False

    return label_check

def get_slice_bboxes_v2(
    image_height: int,
    image_width: int,
    slice_height: int = 480,
    slice_width: int = 640,
    overlap_height_ratio: int = 0.2,
    overlap_width_ratio: int = 0.2,
) -> List[List[int]]:
    """Slices `image_pil` in crops.
    Corner values of each slice will be generated using the `slice_height`,
    `slice_width`, `overlap_height_ratio` and `overlap_width_ratio` arguments.

    Args:
        image_height (int): Height of the original image.
        image_width (int): Width of the original image.
        slice_height (int): Height of each slice. Default 512.
        slice_width (int): Width of each slice. Default 512.
        overlap_height_ratio(float): Fractional overlap in height of each
            slice (e.g. an overlap of 0.2 for a slice of size 100 yields an
            overlap of 20 pixels). Default 0.2.
        overlap_width_ratio(float): Fractional overlap in width of each
            slice (e.g. an overlap of 0.2 for a slice of size 100 yields an
            overlap of 20 pixels). Default 0.2.

    Returns:
        List[List[int]]: List of 4 corner coordinates for each N slices.
            [
                [slice_0_left, slice_0_top, slice_0_right, slice_0_bottom],
                ...
                [slice_N_left, slice_N_top, slice_N_right, slice_N_bottom]
            ]
    """
    slice_bboxes = []
    y_max = y_min = 0
    y_overlap = int(overlap_height_ratio * slice_height)
    x_overlap = int(overlap_width_ratio * slice_width)
    while y_max < image_height:
        x_min = x_max = 0
        y_max = y_min + slice_height
        while x_max < image_width:
            x_max = x_min + slice_width
            if y_max > image_height or x_max > image_width:
                xmax = min(image_width, x_max)
                ymax = min(image_height, y_max)
                xmin = max(0, xmax - slice_width)
                ymin = max(0, ymax - slice_height)
                slice_bboxes.append([xmin, ymin, xmax, ymax])
            else:
                slice_bboxes.append([x_min, y_min, x_max, y_max])
            x_min = x_max - x_overlap
        y_min = y_max - y_overlap
    return slice_bboxes

def get_slice_bboxes(
    image_height: int,
    image_width: int,
    grid_num: int = 2,
    slice_height: int = 480, # no use now
    slice_width: int = 640,
    overlap_height_ratio: int = 0.2,
    overlap_width_ratio: int = 0.2,
) -> List[List[int]]:
    slice_bboxes = []
    half_h = int(image_height // grid_num)
    half_w = int(image_width // grid_num)
    for i in range(grid_num):
        for j in range(grid_num):
            slice_bboxes.append([i * half_h, (i + 1) * half_h, j * half_w, (j + 1) * half_w])
    return slice_bboxes


def main():
    args = parse_args()
    os.makedirs(args.save_dir, exist_ok=True)
    args.no_defect_code = [str(x) for x in args.no_defect_code]
    path_dict = json.load(open(args.json_path,'r'))
    all_train_total = pd.DataFrame()
    all_valid_total = pd.DataFrame()
    random.seed = 752
    for pro_name in path_dict.keys():
        idct_per_pro = path_dict[pro_name]
        for site_name in idct_per_pro.keys():
            idct_per_pro_site = idct_per_pro[site_name]
            for version_name in idct_per_pro_site.keys():
                print(f'Working on {pro_name}-{site_name}-{version_name}.')

                dict_cache = idct_per_pro_site[version_name]
                df_total = pd.DataFrame()
                img_dir = dict_cache['img_dir']
                mask_dir =  dict_cache['mask_dir']
                code_name_list = get_subdirs(img_dir)  # [code for code in get_subdirs(img_dir) if code in args.small_code]
                for code_name in code_name_list:
                    nolabel_count = 0
                    data_dir = os.path.join(img_dir, code_name)
                    # img_ext = get_extensions(data_dir)
                    # img_list = sorted(glob.glob(os.path.join(data_dir, '*.' + str(img_ext)))
                    img_list = []
                    for img_ext in args.img_ext:
                        img_list.extend(sorted(glob.glob(os.path.join(data_dir, '*.' + str(img_ext)))))
                    for img_cache in tqdm(img_list, f'{code_name}'):
                        try:
                            file_ = Image.open(img_cache).load()
                            base_name = os.path.basename(img_cache)
                            img_name = base_name[:-4]
                            code_mask_dir = os.path.join(mask_dir, code_name)
                            maskPath = os.path.join(code_mask_dir, img_name + '.json')
                            mask = None
                            if os.path.exists(maskPath):
                                label_check = check_json(maskPath)
                                if label_check:
                                    mask = json_to_image(img_cache, maskPath, mask_label='box')
                            elif os.path.exists(os.path.join(code_mask_dir, img_name + '.png')):
                                mask = cv2.imread(os.path.join(code_mask_dir, img_name + '.png'), cv2.IMREAD_GRAYSCALE)
                            if mask is not None:
                                code_dist_dir = os.path.join(args.save_dir, code_name)
                                os.makedirs(code_dist_dir, exist_ok=True)
                                img = cv2.imread(img_cache)
                                H, W, _ = img.shape
                                slice_bboxes = get_slice_bboxes(H, W, args.grid_num)
                                for idx, bbox in enumerate(slice_bboxes):
                                    ymin, ymax, xmin, xmax = bbox
                                    if np.sum(mask[ymin:ymax, xmin:xmax]) > 0:
                                        cv2.imwrite('%s/%s_%d.jpg' % (code_dist_dir, img_name, idx), img[ymin:ymax, xmin:xmax, ...])
                                        cv2.imwrite('%s/%s_%d.png' % (code_dist_dir, img_name, idx), mask[ymin:ymax, xmin:xmax] * 255)
                            if code_name in args.no_defect_code:
                                code_dist_dir = os.path.join(args.save_dir, code_name)
                                os.makedirs(code_dist_dir, exist_ok=True)
                                img = cv2.imread(img_cache)
                                H, W, _ = img.shape
                                slice_bboxes = get_slice_bboxes(H, W, args.grid_num)
                                mask = np.zeros(img.shape[:2], dtype=np.int32)
                                for idx, bbox in enumerate(slice_bboxes):
                                    ymin, ymax, xmin, xmax = bbox
                                    # if np.sum(mask[ymin:ymax, xmin:xmax]) > 0:
                                    cv2.imwrite('%s/%s_%d.jpg' % (code_dist_dir, img_name, idx), img[ymin:ymax, xmin:xmax, ...])
                                    cv2.imwrite('%s/%s_%d.png' % (code_dist_dir, img_name, idx), mask[ymin:ymax, xmin:xmax])

                        except IOError:
                            print('Image file %s has something wrong, will not used for train.' % img_cache)



if __name__ == '__main__':
    main()

